import os
import sys
# 确保项目根目录在模块路径里
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
from configs.config import load_config
from data_process.radar_reader import RadarReader

# ── 1. 读取 5000 帧 ──────────────────────────
cfg = load_config("./configs/data_process.py")
cfg['data_batch_size'] = 1000
reader = RadarReader(cfg)
first_root = list(cfg['data_root_list'].values())[0]
radar_data = reader.process_all_files(os.path.join(first_root, 'SARRadar'), 1)

# ── 2. 收集 x, y, cohen ──────────────────────
xs = np.array([pt.x for pts in radar_data.values() for pt in pts])
ys = np.array([pt.y for pts in radar_data.values() for pt in pts])
cohens = np.array([pt.cohen_factor for pts in radar_data.values() for pt in pts])

# ── 3. 2m 分辨率分桶，算平均 ───────────────────
bin_size = 2.0
x_bins = np.arange(xs.min(), xs.max() + bin_size, bin_size)
y_bins = np.arange(ys.min(), ys.max() + bin_size, bin_size)
sum_cohen, _, _ = np.histogram2d(xs, ys, bins=[x_bins, y_bins], weights=cohens)
count,     _, _ = np.histogram2d(xs, ys, bins=[x_bins, y_bins])
with np.errstate(divide='ignore', invalid='ignore'):
    mean_cohen = sum_cohen / count
mean_cohen[count == 0] = np.nan

# ── 4. 计算格子中心 ───────────────────────────
x_centers = (x_bins[:-1] + x_bins[1:]) / 2
y_centers = (y_bins[:-1] + y_bins[1:]) / 2
Xc, Yc = np.meshgrid(x_centers, y_centers, indexing='ij')

# ── 5. 构建对称多项式基，共 10 项 ────────────
#    只用偶次的 x（0,2）和 y^j (j=0..4)
exponents = [(i, j) for i in (0, 2) for j in range(5)]  # 2*5 = 10

# 取有效点用于拟合
valid = ~np.isnan(mean_cohen)
Xv = Xc[valid]
Yv = Yc[valid]
Zv = mean_cohen[valid]

# 设计矩阵 A，shape = (N_samples, 10)
A = np.vstack([ (Xv**i) * (Yv**j) for (i, j) in exponents ]).T

# 最小二乘求解系数
coeffs, *_ = np.linalg.lstsq(A, Zv, rcond=None)

# 拟合函数
def sym_poly10(x, y):
    val = 0.0
    for (i, j), c in zip(exponents, coeffs):
        val += c * (x**i) * (y**j)
    return val

# 在网格上评估拟合曲面
Z_fit = sym_poly10(Xc, Yc)

# ── 6. 掩码无数据区域，避免拟合到零 ────────────
mask = (count == 0)
Z_fit[mask] = np.nan

# ── 7. 绘图对比 ───────────────────────────────
fig = plt.figure(figsize=(12, 6))

# 左：原始分桶曲面
ax1 = fig.add_subplot(121, projection='3d')
surf1 = ax1.plot_surface(
    Xc, Yc, mean_cohen,
    rstride=1, cstride=1, cmap='viridis', edgecolor='none'
)
ax1.set_title('Binned Mean Cohen (2m bins)')
ax1.set_xlabel('X (m)')
ax1.set_ylabel('Y (m)')
ax1.set_zlabel('Mean Cohen')
fig.colorbar(surf1, ax=ax1, shrink=0.5, aspect=10)

# 右：10-参数对称多项式拟合曲面
ax2 = fig.add_subplot(122, projection='3d')
surf2 = ax2.plot_surface(
    Xc, Yc, Z_fit,
    rstride=1, cstride=1, cmap='plasma', edgecolor='none'
)
ax2.set_title('Y‑Axis‑Symmetric 10‑Param Fit')
ax2.set_xlabel('X (m)')
ax2.set_ylabel('Y (m)')
ax2.set_zlabel('Fitted Cohen')
fig.colorbar(surf2, ax=ax2, shrink=0.5, aspect=10)

plt.suptitle('Mean Cohen over XY and Symmetric 10‑Param Fit', fontsize=14)
plt.tight_layout(rect=[0, 0, 1, 0.95])
plt.show()